import cv2
import random
import colorsys
import numpy as np

from tensorflow_core.python.keras.utils import get_file


from core.config import cfg


def load_weights(model, weights_file=None):
    '''
    :param model:
    :param weights_file:
    :return: a list,which contains changed layer's name
    '''
    if weights_file is None:
        weights_file = get_file('yolov3.weights', 'https://pjreddie.com/media/files/yolov3.weights')
    wf = open(weights_file, 'rb')
    major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5)
    changed_layer_names = []
    j = 0
    for i in range(0, 75):
        conv_layer_name = 'conv2d_%d' % i if i > 0 else 'conv2d'
        bn_layer_name = 'batch_normalization_%d' % j if j > 0 else 'batch_normalization'

        conv_layer = model.get_layer(conv_layer_name)

        filters = conv_layer.filters
        k_size = conv_layer.kernel_size[0]
        in_dim = conv_layer.input_shape[-1]
        # print(k_size,fit)

        if i not in [58, 66, 74]:
            # [58,66,74] 代表分类和位置预测层
            # darknet weights: [beta, gamma, mean, variance]
            bn_weights = np.fromfile(wf, dtype=np.float32, count=4 * filters)
            # tf weights: [gamma, beta, mean, variance]
            bn_weights = bn_weights.reshape((4, filters))[[1, 0, 2, 3]]
            bn_layer = model.get_layer(bn_layer_name)
            j += 1
        else:
            conv_bias = np.fromfile(wf, dtype=np.float32, count=filters)

        # darknet shape (out_dim, in_dim, height, width)
        conv_shape = (filters, in_dim, k_size, k_size)
        conv_weights = np.fromfile(wf, dtype=np.float32, count=np.product(conv_shape))
        # tf shape (height, width, in_dim, out_dim)

        conv_weights = conv_weights.reshape(conv_shape).transpose([2, 3, 1, 0])
        if i not in [58, 66, 74]:
            conv_layer.set_weights([conv_weights])
            bn_layer.set_weights(bn_weights)
            changed_layer_names.append(bn_layer.name)
        else:
            conv_layer.set_weights([conv_weights, conv_bias])
            changed_layer_names.append(conv_layer.name)
    assert len(wf.read()) == 0, 'failed to read all data'
    wf.close()
    return changed_layer_names


def read_class_names(class_file_name):
    '''loads class name from a file'''
    names = {}
    with open(class_file_name, 'r') as data:
        for ID, name in enumerate(data):
            names[ID] = name.strip('\n')
    return names


def get_anchors(anchors_path):
    '''loads the anchors from a file'''
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = np.array(anchors.split(','), dtype=np.float32)
    return anchors.reshape(3, 3, 2)


def image_preporcess(image, target_size, gt_boxes=None):
    ih, iw = target_size
    h, w, _ = image.shape

    scale = min(iw / w, ih / h)
    nw, nh = int(scale * w), int(scale * h)
    image_resized = cv2.resize(image, (nw, nh))

    image_paded = np.full(shape=[ih, iw, 3], fill_value=128.0)
    dw, dh = (iw - nw) // 2, (ih - nh) // 2
    image_paded[dh:nh + dh, dw:nw + dw, :] = image_resized
    image_paded = image_paded / 255.

    if gt_boxes is None:
        return image_paded

    else:
        gt_boxes[:, [0, 2]] = gt_boxes[:, [0, 2]] * scale + dw
        gt_boxes[:, [1, 3]] = gt_boxes[:, [1, 3]] * scale + dh
        return image_paded, gt_boxes

def relu(x):
    if x > 0:
        return x
    else:
        return 0

def draw_bbox(image, bboxes, pose_score_threshold=-5, classes=None, person_poses_classes=None, show_label=True, ):
    """
    bboxes: [x_min, y_min, x_max, y_max, probability, cls_id] format coordinates.
    """

    if classes is None:
        classes = read_class_names(cfg.YOLO.CLASSES)
    if person_poses_classes is None:
        person_poses_classes = read_class_names(cfg.YOLO.PERSON_POSES_CLASSES)
    num_color_types = len(classes) + len(person_poses_classes)

    image_h, image_w, _ = image.shape
    hsv_tuples = [(1.0 * x / num_color_types, 1., 1.) for x in range(num_color_types)]
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))

    random.seed(4)
    random.shuffle(colors)
    random.seed(None)

    class_names = {}
    for i, bbox in enumerate(bboxes):
        coor = np.array(bbox[:4], dtype=np.int32)
        fontScale = min(image_h,image_w) * 0.0015
        score = bbox[4]
        class_ind = int(bbox[5])
        # pose_ind = int(bboxes[])
        bbox_color = colors[class_ind]
        poses_ind = int(bbox[6])
        poses_score = bbox[7]
        poses_color = colors[class_ind+poses_ind]

        bbox_thick = int(0.6 * (image_h + image_w) / 800)
        c1, c2 = (coor[0], coor[1]), (coor[2], coor[3])
        cv2.rectangle(image, c1, c2, bbox_color, bbox_thick)

        class_names[classes[class_ind]] = bbox_color

        if show_label:
            bbox_mess = '%s: %.2f' % (classes[class_ind], score)
            if poses_score > pose_score_threshold:
                poses_mess = '%s: %.2f' % (person_poses_classes[poses_ind], poses_score)
            else:
                poses_mess = 'unknown'
            t_size1 = cv2.getTextSize(bbox_mess, 0, fontScale, thickness=bbox_thick // 2)[0]
            t_size2 = cv2.getTextSize(poses_mess, 0, fontScale, thickness=bbox_thick // 2)[0]
            box_length = max(t_size1[0],t_size2[0])
            brightness = 50
            coor1 = (c1[0] + 1, c2[1]-int(t_size1[1]*1.5))
            coor2 = (coor1[0] + box_length , coor1[1] + int(t_size1[1]*1.5))
            coor3 = (coor1[0] , coor1[1] + t_size1[1] + 1)
            cv2.rectangle(image, coor1, coor2, bbox_color, -1)  # filled
            cv2.putText(image, bbox_mess, coor3, cv2.FONT_HERSHEY_SIMPLEX,
                        fontScale, (brightness-bbox_color[0], brightness-bbox_color[1], brightness-bbox_color[2]), bbox_thick // 2, lineType=cv2.LINE_AA)


            if classes[class_ind] == 'person':
                coor1_ = (coor1[0], coor2[1]-3*t_size2[1])
                coor2_ = (coor1_[0] + box_length, coor1_[1] + int(t_size2[1] *1.5))
                coor3_ = (coor1_[0], coor1_[1] + t_size2[1] + 1)
                cv2.rectangle(image, coor1_, coor2_, poses_color, -1)  # filled
                cv2.putText(image, poses_mess, coor3_, cv2.FONT_HERSHEY_SIMPLEX,
                            fontScale, (relu(brightness-poses_color[0]), relu(brightness-poses_color[1]), relu(brightness-poses_color[2])),
                            bbox_thick // 2, lineType=cv2.LINE_AA)

    print(class_names)
    return image


def bboxes_iou(boxes1, boxes2):
    boxes1 = np.array(boxes1)
    boxes2 = np.array(boxes2)

    boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1])
    boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1])

    left_up = np.maximum(boxes1[..., :2], boxes2[..., :2])
    right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:])

    inter_section = np.maximum(right_down - left_up, 0.0)
    inter_area = inter_section[..., 0] * inter_section[..., 1]
    union_area = boxes1_area + boxes2_area - inter_area
    ious = np.maximum(1.0 * inter_area / union_area, np.finfo(np.float32).eps)

    return ious


def nms(bboxes, iou_threshold, sigma=0.3, method='nms'):
    """
    :param bboxes: (xmin, ymin, xmax, ymax, score, class)

    Note: soft-nms, https://arxiv.org/pdf/1704.04503.pdf
          https://github.com/bharatsingh430/soft-nms
    """
    classes_in_img = list(set(bboxes[:, 5]))
    best_bboxes = []

    for cls in classes_in_img:
        cls_mask = (bboxes[:, 5] == cls)
        cls_bboxes = bboxes[cls_mask]

        while len(cls_bboxes) > 0:
            max_ind = np.argmax(cls_bboxes[:, 4])
            best_bbox = cls_bboxes[max_ind]
            best_bboxes.append(best_bbox)
            cls_bboxes = np.concatenate([cls_bboxes[: max_ind], cls_bboxes[max_ind + 1:]])
            iou = bboxes_iou(best_bbox[np.newaxis, :4], cls_bboxes[:, :4])
            weight = np.ones((len(iou),), dtype=np.float32)

            assert method in ['nms', 'soft-nms']

            if method == 'nms':
                iou_mask = iou > iou_threshold
                weight[iou_mask] = 0.0

            if method == 'soft-nms':
                weight = np.exp(-(1.0 * iou ** 2 / sigma))

            cls_bboxes[:, 4] = cls_bboxes[:, 4] * weight
            score_mask = cls_bboxes[:, 4] > 0.
            cls_bboxes = cls_bboxes[score_mask]

    return best_bboxes


def postprocess_boxes(pred_bbox, org_img_shape, input_size, class_score_threshold):
    NUM_CLASS = len(read_class_names(cfg.YOLO.CLASSES))
    # NUM_POSES = len(read_class_names(cfg.YOLO.PERSON_POSES_CLASSES))
    valid_scale = [0, np.inf]
    pred_bbox = np.array(pred_bbox)

    pred_xywh = pred_bbox[:, 0:4]
    pred_conf = pred_bbox[:, 4]
    pred_prob = pred_bbox[:, 5:5 + NUM_CLASS]
    pred_person_prob = pred_bbox[:, 5]
    pose_prob = pred_bbox[:, 5 + NUM_CLASS:]

    # # (1) (x, y, w, h) --> (xmin, ymin, xmax, ymax)
    pred_coor = np.concatenate([pred_xywh[:, :2] - pred_xywh[:, 2:] * 0.5,
                                pred_xywh[:, :2] + pred_xywh[:, 2:] * 0.5], axis=-1)
    # # (2) (xmin, ymin, xmax, ymax) -> (xmin_org, ymin_org, xmax_org, ymax_org)
    org_h, org_w = org_img_shape
    resize_ratio = min(input_size / org_w, input_size / org_h)

    dw = (input_size - resize_ratio * org_w) / 2
    dh = (input_size - resize_ratio * org_h) / 2

    pred_coor[:, 0::2] = 1.0 * (pred_coor[:, 0::2] - dw) / resize_ratio
    pred_coor[:, 1::2] = 1.0 * (pred_coor[:, 1::2] - dh) / resize_ratio

    # # (3) clip some boxes those are out of range
    pred_coor = np.concatenate([np.maximum(pred_coor[:, :2], [0, 0]),
                                np.minimum(pred_coor[:, 2:], [org_w - 1, org_h - 1])], axis=-1)
    invalid_mask = np.logical_or((pred_coor[:, 0] > pred_coor[:, 2]), (pred_coor[:, 1] > pred_coor[:, 3]))
    pred_coor[invalid_mask] = 0

    # # (4) discard some invalid boxes
    bboxes_scale = np.sqrt(np.multiply.reduce(pred_coor[:, 2:4] - pred_coor[:, 0:2], axis=-1))
    scale_mask = np.logical_and((valid_scale[0] < bboxes_scale), (bboxes_scale < valid_scale[1]))

    # # (5) discard some boxes with low scores
    classes = np.argmax(pred_prob, axis=-1)

    scores = pred_conf * pred_prob[np.arange(len(pred_coor)), classes]
    # print(score)
    poses_classes = np.argmax(pose_prob, axis=-1)
    poses_socres = (pred_person_prob * pose_prob[np.arange(len(pred_coor)), poses_classes])
    score_mask = scores > class_score_threshold

    mask = np.logical_and(scale_mask, score_mask)
    coors, scores, classes, poses_classes, poses_socres = pred_coor[mask], scores[mask], classes[mask], \
                                                          poses_classes[mask], poses_socres[mask]

    return np.concatenate([coors, scores[:, np.newaxis], classes[:, np.newaxis], poses_classes[:, np.newaxis],
                           poses_socres[:, np.newaxis]], axis=-1)
